import os
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
import base64
from conf import conf

def load_rsa_key():
    base_dir = os.path.dirname(os.path.abspath(__file__)).replace('\\', '/') + "/../../"
    key_path = base_dir + "data/server/"
    pk1,pk2 = None,None
    with open(key_path + "pk1.pem", "rb") as key_file:
        pk1 = RSA.import_key(key_file.read())
    with open(key_path + "pk2.pem", "rb") as key_file:
        pk2 = RSA.import_key(key_file.read())
    return pk1,pk2



def work():
    config = conf.read_json()
    ks =  bytes.fromhex(config.get('select_key'))
    pk1,pk2 = load_rsa_key()
    cipher1 = PKCS1_OAEP.new(pk1) #0
    cipher2 = PKCS1_OAEP.new(pk2) #1
    base_dir = os.path.dirname(os.path.abspath(__file__)).replace('\\', '/') + "/../../data/server/"
    sel = []
    for i, b in enumerate(ks):
        for j in range(8):
            sel.append((b >> (7 - j)) & 0x01)

    with open(base_dir+"/key.txt", "r") as fin,  open(base_dir+"/key_enc.txt", "w")as fout:
        for lineno, line in enumerate(fin, start=0):
            hex_string = line.strip()
            if not hex_string:
                continue  # 跳过空行
            try:
                # 把16进制字符串转换为bytes（长度应该是16字节）
                data = bytes.fromhex(hex_string.strip())
                if len(data) != 16:
                    raise ValueError("行数据不是128位（16字节）")
                if sel[lineno] == 0:
                    data = cipher1.encrypt(data)
                else:
                    data = cipher2.encrypt(data)
                encrypted_str = base64.b64encode(data).decode()
                fout.write(encrypted_str + "\n")
            except Exception as e:
                print(f"第 {lineno} 行加密失败: {e}")


if __name__ == '__main__':
    work()